import pandas as pd
from scipy.stats import chi2_contingency

from src.zhangyu import label_encode

data = pd.read_csv('../../data/raw/train.csv')
# 数据预处理  将有多分类的列 转化为 数值型的
# 需要编码的列
cols = [
    "BusinessTravel", "Department", "Education", "EducationField",
    "EnvironmentSatisfaction", "Gender", "JobInvolvement", "JobLevel", "JobRole",
    "JobSatisfaction", "MaritalStatus", "OverTime", "PerformanceRating",
    "RelationshipSatisfaction", "TrainingTimesLastYear", "StockOptionLevel", "WorkLifeBalance"
]
data = label_encode.encode(data, cols)
categorical_cols = ["BusinessTravel", "Department", "Education", "EducationField", "EnvironmentSatisfaction",
                    "Gender", "JobInvolvement", "JobLevel", "JobRole",
                    "JobSatisfaction", "MaritalStatus", "OverTime", "PerformanceRating",
                    "RelationshipSatisfaction", "TrainingTimesLastYear", "StockOptionLevel", "WorkLifeBalance"]
print("类别型特征：", categorical_cols)
target = 'Attrition'
chi2_results = []
for feature in categorical_cols:
    print(f"\n=== 卡方检验: {feature} vs {target} ===")

    # 构建列联表
    contingency = pd.crosstab(data[feature], data[target])

    # 执行卡方检验
    chi2, p, dof, expected = chi2_contingency(contingency)

    # 输出结果
    # print(f"Chi-square: {chi2:.4f}")
    # print(f"P-value: {p:.4f}")
    # print(f"自由度: {dof}")
    #
    # if p < 0.05:
    #     print(f"→ {feature} 与 {target} 存在显著关联")
    # else:
    #     print(f"→ {feature} 与 {target} 没有显著关联")
    chi2_results.append({
        'Feature': feature,
        'Chi-square': chi2,
        'P-value': p,
        'DOF': dof,
        'Significant': p < 0.05
    })

    # 将结果转为 DataFrame 并按 p 值排序
pd.set_option('display.float_format', lambda x: '%.6f' % x)

# 输出排序后的结果

results_df = pd.DataFrame(chi2_results)
results_df = results_df.sort_values(by='Significant', ascending=True)

# 输出排序后的结果
print(results_df)
